import fasttext
import pandas as pd
from sklearn.metrics import f1_score
from chat_bot.models.intent_fasttext import config

fasttext.FastText.eprint = lambda x: None


# 十折交叉验证
def testFasttext():
    fscore_list = []
    for i in range(10):
        # 拿数据
        train_path = './data/train' + str(i) + '.csv'
        test_path = './data/test' + str(i) + '.csv'
        test_df = pd.read_csv(test_path, encoding='utf-8', sep="\t")
        model = fasttext.train_supervised(train_path, lr=config.lr, wordNgrams=config.word_ngrams, dim=config.dim,
                                          epoch=config.epoch,
                                          loss='softmax')
        test_pred = [model.predict(x)[0][0] for x in test_df['text'].values.astype(str)]
        f1 = f1_score(test_df['label'].values.astype(str), test_pred, average='macro')
        fscore_list.append(f1)
        print(f"the f1_score of training{i + 1} is: ", f1)
    print('The average f1_score is', sum(fscore_list) / len(fscore_list))


def printResults(N, p, r):
    # 数据量
    print("N\t" + str(N))
    # 精确率
    print("P@{}\t{:.3f}".format(1, p))
    # 召回率
    print("R@{}\t{:.3f}".format(1, r))


# 模型训练和保存
def trainFasttext():
    # model = fasttext.train_supervised('data/train_true.csv', autotuneValidationFile='data/test.csv')
    model = fasttext.train_supervised(config.all_data_path, lr=config.lr, dim=config.dim, epoch=config.epoch,
                                      word_ngrams=config.word_ngrams, loss='softmax')
    # 保存模型
    model.save_model(config.model_path)


if __name__ == "__main__":
    testFasttext()
    trainFasttext()
    print('模型构建成功')

    # 测试
    # model = fasttext.load_model(config.model_path)
    # print(model.test_label(config.test_data_path, k=1))
    # printResults(*model.test(config.test_data_path))
